package kafka

import (
	"context"
	"github.com/Shopify/sarama"
)

type KafkaProducer struct {
	sarama.SyncProducer
}
type KafkaConsumer struct {
	sarama.Consumer
	OffsetMode int64
}

// NewKafkaProducer 创建一个kafka生产者
func NewKafkaProducer(info *KafkaConfig) (*KafkaProducer, error) {
	config := sarama.NewConfig()
	config.Producer.RequiredAcks = sarama.RequiredAcks(info.RequiredAckType) // 发送完数据需要leader和follow都确认 NoResponse = 0 ,WaitForLocal = 1 ,WaitForAll = -1
	config.Producer.Partitioner = sarama.NewRandomPartitioner                // 新选出一个partition
	config.Producer.Return.Successes = true                                  // 成功交付的消息将在success channel返回

	producer, err := sarama.NewSyncProducer(info.Endpoints, config)
	if err != nil {
		return nil, err
	} else {
		return &KafkaProducer{producer}, nil
	}
}

// Close 关闭生产者
func (p *KafkaProducer) Close() error {
	return p.SyncProducer.Close()
}

// SendMessage 发生一个消息
func (p *KafkaProducer) SendMessage(topic, massage string) (partition int32, offset int64, err error) {
	// 构造一个消息
	msg := &sarama.ProducerMessage{}
	msg.Topic = topic
	msg.Value = sarama.StringEncoder(massage)

	// 发送消息
	return p.SyncProducer.SendMessage(msg)
}

// SendMessages 批量发生消息
func (p *KafkaProducer) SendMessages(messages map[string][]string) (err error) {
	if messages != nil && len(messages) > 0 {
		msgs := make([]*sarama.ProducerMessage, 0)
		for k, v := range messages {
			for _, m := range v {
				// 构造一个消息
				msg := &sarama.ProducerMessage{}
				msg.Topic = k
				msg.Value = sarama.StringEncoder(m)
				msgs = append(msgs, msg)
			}
		}
		// 发送消息
		return p.SyncProducer.SendMessages(msgs)
	}
	return nil
}

// NewKafkaConsumer 生产一个消费者
func NewKafkaConsumer(info *KafkaConfig) (*KafkaConsumer, error) {
	consumer, err := sarama.NewConsumer(info.Endpoints, nil)
	if err != nil {
		return nil, err
	} else {
		return &KafkaConsumer{consumer, info.OffsetMode}, nil
	}

}

// Close 关闭消费者
func (c *KafkaConsumer) Close() error {
	return c.Consumer.Close()
}

// Consume 消费者消费某个主题  调用该方法回阻塞
func (c *KafkaConsumer) Consume(ctx context.Context, topic string, callback func(partition int32, offset int64, key, value string) error) error {
	partitionList, err := c.Partitions(topic) // 根据topic取到所有的分区
	if err != nil {
		return err
	}
	for partition := range partitionList { // 遍历所有的分区
		// 针对每个分区创建一个对应的分区消费者
		pc, err := c.ConsumePartition(topic, int32(partition), c.OffsetMode) // OffsetNewest -1;OffsetOldest -2
		if err != nil {
			return err
		}
		defer pc.AsyncClose()
		// 异步从每个分区消费信息
		go func(sarama.PartitionConsumer) {
			for msg := range pc.Messages() {
				callback(msg.Partition, msg.Offset, string(msg.Key), string(msg.Value))
				//fmt.Printf("Partition:%d Offset:%d Key:%v Value:%v", msg.Partition, msg.Offset, msg.Key, msg.Value)
			}
		}(pc)
	}
	<-ctx.Done()
	return nil
}
